from pydantic import BaseModel
from typing import List, Optional, Union


# Модель для ответа на запрос по получению количества токенов - https://developers.sber.ru/docs/ru/gigachat/api/reference/rest/post-tokens-count

class TokensCountResponseModel(BaseModel):
    tokens: int
    characters: int


# Модель для ответа на запрос по получению доступных моделей - https://developers.sber.ru/docs/ru/gigachat/api/reference/rest/get-models

class ModelsResponseModel(BaseModel):


    class Data(BaseModel):
        id: str
        object: str
        owned_by: str
    
    data: List[Data]
    object: str

# Модель для ответа на запрос модели - https://developers.sber.ru/docs/ru/gigachat/api/reference/rest/post-chat

class CompletionResponseModel(BaseModel):


    class Choice(BaseModel):

        class Message(BaseModel):
            role:str
            content:str
            data_for_context: List

        message: Message
        index: int
        finish_reason: str
    
    

    class Usage(BaseModel):
        prompt_tokens: int
        completion_tokens: int
        total_tokens: int
        system_tokens: int
    
    choices: List[Choice]
    created: int
    model: str
    usage: Usage
    object: str

#модель ответа для метода - https://developers.sber.ru/docs/ru/gigachat/api/reference/rest/post-embeddings
class CreateEmbeddingResponseModel(BaseModel):
    
    class Data(BaseModel):
        class Usage(BaseModel):
            prompt_tokens: int
    

        object: str
        embedding: List[float]
        index: int
        usage: Usage


    object: str
    data: List[Data]
    model: Optional[str]


#